#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.executors import ExternalShutdownException
from rclpy.qos import qos_profile_sensor_data
from std_msgs.msg import Bool
from geometry_msgs.msg import PoseStamped, TransformStamped
from rclpy.action import ActionClient
from moveit_msgs.action import MoveGroup
from moveit_msgs.msg import MoveItErrorCodes
from tf2_ros import Buffer, TransformListener, TransformBroadcaster
from tf2_ros.transformations import quaternion_matrix
import numpy as np
import copy

# 全局变量
right_pose = None
left_pose = None
last_valid_pose = None
task = "press"
start_move = False  # 初始时，start_move 设为 False

# 回调函数，用于接收 right_pose
def right_pose_callback(msg):
    global right_pose
    right_pose = msg

# 回调函数，用于接收 left_pose
def left_pose_callback(msg):
    global left_pose, last_valid_pose, task, start_move

    if not start_move:  # 如果 start_move 为 False，则不执行动作
        print("Waiting for start_move signal to be True.")
        return

    # 检查当前 pose 是否有效（避免接收到无效位姿）
    if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
        if last_valid_pose:
            print("Invalid pose received, using last valid pose.")
            left_pose = copy.deepcopy(last_valid_pose)
        else:
            print("No valid pose available, skipping this callback.")
            return
    else:
        # 计算当前位置与上次有效位置之间的距离
        if last_valid_pose:
            current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
            last_position = np.array([last_valid_pose.pose.position.x, last_valid_pose.pose.position.y, last_valid_pose.pose.position.z])
            distance = np.linalg.norm(current_position - last_position)

            # 如果距离大于 1cm，视为无效位姿，不进行规划
            if distance > 0.01:
                print(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                return
            else:
                left_pose = copy.deepcopy(msg)
                last_valid_pose = copy.deepcopy(msg)  # 保存当前有效的 pose
        else:
            left_pose = copy.deepcopy(msg)
            last_valid_pose = copy.deepcopy(msg)  # 保存当前有效的 pose

    print("Received a pose, starting movement.")

    quaternion = [left_pose.pose.orientation.x, left_pose.pose.orientation.y, left_pose.pose.orientation.z, left_pose.pose.orientation.w]

    br = TransformBroadcaster()
    t = TransformStamped()
    t.header.stamp = rclpy.time.Time().to_msg()
    t.header.frame_id = "base_link"
    t.child_frame_id = "button_pose"

    t.transform.translation.x = left_pose.pose.position.x
    t.transform.translation.y = left_pose.pose.position.y
    t.transform.translation.z = left_pose.pose.position.z

    t.transform.rotation.x = left_pose.pose.orientation.x
    t.transform.rotation.y = left_pose.pose.orientation.y
    t.transform.rotation.z = left_pose.pose.orientation.z
    t.transform.rotation.w = left_pose.pose.orientation.w
    br.sendTransform(t)

    rotation_matrix = quaternion_matrix(quaternion)
    
    local_offset1 = np.array([-0.01, 0.0, 0.00, 1.0])  # [世界坐标系y轴，世界坐标系z轴，世界坐标系x轴]
    local_offset2 = np.array([-0.01, 0.0, 0.05, 1.0])  # [世界坐标系y轴，世界坐标系z轴，世界坐标系x轴]
    world_offset1 = np.dot(rotation_matrix, local_offset1)
    world_offset2 = np.dot(rotation_matrix, local_offset2)

    if task == "press":
        print(task)
        msg_hand_control1 = Bool()
        msg_hand_control1.data = True

        left_pose.pose.position.x += world_offset1[0]
        left_pose.pose.position.y += world_offset1[1]
        left_pose.pose.position.z += world_offset1[2]

        left_pose.pose.orientation.x = -0.10711095938856535
        left_pose.pose.orientation.y = -0.33227448354515043
        left_pose.pose.orientation.z = -0.845326815249489
        left_pose.pose.orientation.w = 0.40440510059075774

        move_to_pose(left_pose, "Left_Arm_Link7")
        start_move = False
        task = "release"

    elif task == "release":
        print(task)
        msg_hand_control1 = Bool()
        msg_hand_control1.data = True

        left_pose.pose.position.x += world_offset2[0]
        left_pose.pose.position.y += world_offset2[1]
        left_pose.pose.position.z += world_offset2[2]

        left_pose.pose.orientation.x = -0.10711095938856535
        left_pose.pose.orientation.y = -0.33227448354515043
        left_pose.pose.orientation.z = -0.845326815249489
        left_pose.pose.orientation.w = 0.40440510059075774

        move_to_pose(left_pose, "Left_Arm_Link7")

        start_move = False
        task = "press"

# 回调函数，用于接收 start_move 信号
def start_move_callback(msg):
    global start_move
    if msg.data:
        print("Received start_move signal, proceeding with motion.")
        start_move = True
    else:
        print("Received stop_move signal, halting motion.")
        start_move = False

# 移动手臂到指定位姿
def move_to_pose(pose, link_name):
    global move_group_client
    print("adjusted_left_pose", pose)
    print(move_group_client.get_current_pose().pose)

    current_pose = move_group_client.get_current_pose().pose

    # 计算当前位置与目标位置之间的差距
    distance = np.linalg.norm(np.array([pose.pose.position.x, pose.pose.position.y, pose.pose.position.z]) -
                              np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z]))
    print(f"Distance to target pose: {distance}")

    # 设定一个阈值，如果目标位姿与当前位置的距离大于这个值，则进行分步规划
    step_size = 0.01  # 每一步的距离，单位米
    if distance > step_size:
        num_steps = int(np.ceil(distance / step_size))  # 计算需要的步数

        # 按照分步的策略逐步移动
        for i in range(num_steps):
            # 计算每一步的目标位置
            interpolation_factor = (i + 1) / num_steps  # 计算每一步的插值因子
            intermediate_pose = PoseStamped()
            intermediate_pose.header.frame_id = "base_link"
            intermediate_pose.pose.position.x = current_pose.position.x + (pose.pose.position.x - current_pose.position.x) * interpolation_factor
            intermediate_pose.pose.position.y = current_pose.position.y + (pose.pose.position.y - current_pose.position.y) * interpolation_factor
            intermediate_pose.pose.position.z = current_pose.position.z + (pose.pose.position.z - current_pose.position.z) * interpolation_factor

            # 使用当前的位姿插值更新目标姿态
            intermediate_pose.pose.orientation = pose.pose.orientation  # 保持目标的旋转信息不变

            move_group_client.set_pose_target(intermediate_pose)
            plan, traj, planning_time, error_code = move_group_client.plan()

            if plan:
                success = move_group_client.execute(traj, wait=True)
                if not success:
                    print(f"Failed to move {link_name} to intermediate pose at step {i+1}.")
                    return  # 如果某一步执行失败，则终止
            else:
                print(f"Failed to generate plan for {link_name} at step {i+1}.")
                return  # 如果某一步规划失败，则终止

        print(f"Successfully moved {link_name} to the target pose in {num_steps} steps.")
    else:
        # 如果距离较小，直接规划到目标位姿
        move_group_client.set_pose_target(pose)
        plan, traj, planning_time, error_code = move_group_client.plan()

        if plan:
            success = move_group_client.execute(traj, wait=True)
            if success:
                print(f"Successfully moved {link_name} to the target pose.")
            else:
                print(f"Failed to move {link_name} to the target pose.")
        else:
            print(f"Failed to generate plan for {link_name}.")

class MoveArmNode(Node):
    def __init__(self):
        super().__init__('move_arm_to_poses')
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer)
        self.tf_broadcaster = TransformBroadcaster(self)

        self.hand_control_publisher = self.create_publisher(Bool, "/hand_control_topic", qos_profile_sensor_data)
        self.create_subscription(PoseStamped, '/button_pose', left_pose_callback, qos_profile_sensor_data)
        self.create_subscription(Bool, '/start_move', start_move_callback, qos_profile_sensor_data)

        self.move_group_client = ActionClient(self, MoveGroup, 'move_group')

    def __del__(self):
        self.move_group_client.destroy()

def main(args=None):
    rclpy.init(args=args)
    try:
        node = MoveArmNode()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        rclpy.try_shutdown()
        print("Program exited.")

if __name__ == "__main__":
    main()
